#include "pose_observe.h"
#include "ceres_solver_optimize.h"

// obs = [px, py, vx, vy, theta];

// predict - obs
void PoseObserve::ComputeResidualAndJacobian(const State& state, const Eigen::Vector3d &pose, const Eigen::Vector2d &velocity, Eigen::Matrix<double, 5, 1> &residual, Eigen::Matrix<double, 5, 8> &jacobian) {
	Eigen::Vector2d res_p = state.p - pose.block<2,1>(0,0);
	// double res_theta = math_utils::normalize_angle(state.theta - pose[2]);
	double res_theta = math_utils::theta_normalize(state.theta - pose[2]);

	Eigen::Vector2d res_v = state.v - velocity;
  	
  	// cout << "state.v ================== : " << state.v.transpose() << endl;
  	// cout << "velocity ================== : " << velocity.transpose() << endl;
  	// cout << "state.p ================== : " << state.p.transpose() << endl;
  	// cout << "pose ================== : " << pose.block<2,1>(0,0).transpose() << endl;
  	// cout << "state.theta ================== : " << state.theta << endl;
  	// cout << "pose[2] ================== : " << pose[2] << endl;

	residual.block<2,1>(0,0) = res_p;
	residual.block<2,1>(2,0) = res_v;
	residual(4,0) = res_theta;
	
	jacobian = Eigen::Matrix<double, 5, 8>::Zero();
	jacobian.block<5,5>(0,0) = Eigen::Matrix<double, 5, 5>::Identity();

}

void PoseObserve::ComputeResidualAndJacobian(const State& state, const Eigen::Vector3d &pose, const Eigen::Vector2d &velocity, Eigen::Matrix<double, 2, 1> &residual, Eigen::Matrix<double, 2, 8> &jacobian) {
	Eigen::Vector2d res_p = state.p - pose.block<2,1>(0,0);

	residual = res_p;

	jacobian = Eigen::Matrix<double, 2, 8>::Zero();
	jacobian.block<2,2>(0,0) = Eigen::Matrix<double, 2, 2>::Identity();

}

void PoseObserve::ComputeResidualAndJacobianPoseGraph(const State& last_state, const Eigen::Vector3d &last_pose, const Eigen::Vector2d &last_velocity, const State& state, const Eigen::Vector3d &pose, const Eigen::Vector2d &velocity, Eigen::Matrix<double, 5, 1> &residualXi, Eigen::Matrix<double, 5, 8> &jacobianXi, Eigen::Matrix<double, 5, 1> &residualXj, Eigen::Matrix<double, 5, 5> &jacobianXj, Eigen::Matrix<double, 5, 5> &jacobianXij, Eigen::Matrix<double, 5, 8> &jacobianXji) {
	PoseObserve::ComputeResidual(last_state, last_pose, last_velocity, residualXi);
	PoseObserve::ComputeResidual(state, pose, velocity, residualXj);

	// cout << "residualXi ================== : " << residualXi << endl;

	jacobianXi = Eigen::Matrix<double, 5, 8>::Zero();
	jacobianXi.block<5,5>(0,0) = Eigen::Matrix<double, 5, 5>::Identity();
	
    // cout << "jacobianXi == : \n" << jacobianXi << endl;

	jacobianXij = Eigen::Matrix<double, 5, 5>::Zero();

	Eigen::Matrix<double,2,2> R = math_utils::rotation(last_state.theta);

    Eigen::Vector3d last_imu_unbias(state.last_imu[0], state.last_imu[1], state.last_imu[2]);
    Eigen::Vector3d curr_imu_unbias(state.curr_imu[0], state.curr_imu[1], state.curr_imu[2]);
 
    last_imu_unbias[0] -= state.ba[0];
    curr_imu_unbias[0] -= state.ba[0];
    last_imu_unbias[1] -= state.ba[1];
    curr_imu_unbias[1] -= state.ba[1];
    last_imu_unbias[2] -= state.bg;
    curr_imu_unbias[2] -= state.bg;
    
    // std::cout << "dt ============= : " << dt << ", " << init_bias_ << std::endl;
    Eigen::Vector3d imu_unbias_hat = 0.5 * (last_imu_unbias + curr_imu_unbias);
    Eigen::Matrix<double,2,2> sign = Eigen::Matrix2d::Zero();
    sign(0,1) = -1;
    sign(1,0) = 1; 
    Eigen::Matrix<double,2,2> R_dev = sign * R;
	jacobianXji = Eigen::Matrix<double, 5, 8>::Zero();
	jacobianXji.block<2, 2>(0, 0) = Eigen::Matrix<double, 2, 2>::Identity();
	jacobianXji.block<2, 2>(0, 2) = state.dt * Eigen::Matrix<double, 2, 2>::Identity();
	jacobianXji.block<2, 1>(0, 4) = 0.5 * state.dt * state.dt * R_dev * imu_unbias_hat.head<2>();
	jacobianXji.block<2, 2>(0, 5) = 0.5 * state.dt * state.dt * R;

	jacobianXji.block<2, 2>(2, 2) = Eigen::Matrix<double, 2, 2>::Identity();
	jacobianXji.block<2, 1>(2, 4) = 0.5 * state.dt * R_dev * imu_unbias_hat.head<2>();
	jacobianXji.block<2, 2>(2, 5) = state.dt * R;

	jacobianXji(4, 7) = state.dt;

	jacobianXj = Eigen::Matrix<double, 5, 5>::Identity();
}

Eigen::Matrix<double, 10, 1> ComputeResidualVector(const Eigen::Matrix<double, 13, 1> &state, const Eigen::Vector3d &last_pose, const Eigen::Vector2d &last_velocity, const Eigen::Vector3d &pose, const Eigen::Vector2d &velocity) {
	Eigen::Vector2d res_p = state.head<2>() - last_pose.block<2,1>(0,0);
	// double res_theta = math_utils::normalize_angle(state(4,0) - last_pose[2]);
	double res_theta = math_utils::theta_normalize(state(4,0) - last_pose[2]);

	Eigen::Vector2d res_v = state.block<2,1>(2,0) - last_velocity;

	Eigen::Matrix<double, 10, 1> residual;
	residual.block<2,1>(0,0) = res_p;
	residual.block<2,1>(2,0) = res_v;
	residual(4,0) = res_theta;	

	residual.block<2,1>(5,0) = state.block<2,1>(8,0) - pose.block<2,1>(0,0);
	residual.block<2,1>(7,0) = state.block<2,1>(10,0) - velocity.block<2,1>(0,0);
	residual(9,0) = math_utils::theta_normalize(state(12,0) - pose[2]);	

	return residual;
}

void PoseObserve::ComputeResidualAndNumericalJacobian(const State& last_state, const Eigen::Vector3d &last_pose, const Eigen::Vector2d &last_velocity, const State& state, const Eigen::Vector3d &pose, const Eigen::Vector2d &velocity, Eigen::Matrix<double, 5, 1> &residualXi, Eigen::Matrix<double, 5, 8> &jacobianXi, Eigen::Matrix<double, 5, 1> &residualXj, Eigen::Matrix<double, 5, 5> &jacobianXj, Eigen::Matrix<double, 5, 5> &jacobianXij, Eigen::Matrix<double, 5, 8> &jacobianXji) {
	PoseObserve::ComputeResidual(last_state, last_pose, last_velocity, residualXi);
	PoseObserve::ComputeResidual(state, pose, velocity, residualXj);
	Eigen::Matrix<double, 13, 1> state_tmp;
	state_tmp.block<2,1>(0,0) = last_state.p;
	state_tmp.block<2,1>(2,0) = last_state.v;
	state_tmp(4,0) = last_state.theta;

	state_tmp.block<2,1>(5,0) = last_state.ba;
	state_tmp(7,0) = last_state.bg;
	state_tmp.block<2,1>(8,0) = state.p;
	state_tmp.block<2,1>(10,0) = state.v;
	state_tmp(12,0) = state.theta;

	// std::cout << " residualXi : \n" << residualXi << std::endl;
	// std::cout << " residualXj : \n" << residualXj << std::endl;

	const double eps = 1.0e-6;
	const double scalar = 1.0 / (2.0*eps);	
    Eigen::Matrix<double,10,13> jacobian = Eigen::Matrix<double,10,13>::Zero();

	// Eigen::Matrix<double, 10, 1> residual_tmp = ComputeResidualVector(state_tmp, last_pose, last_velocity, pose, velocity);
	// std::cout << " residual_tmp : \n" << residual_tmp << std::endl;

	for(int i = 0; i < 13; i++) {
		Eigen::Matrix<double, 13, 1> state_plus = state_tmp;
		Eigen::Matrix<double, 13, 1> state_minus = state_tmp;

		state_plus[i] += eps;
		state_minus[i] -= eps;
		Eigen::Matrix<double, 10, 1> residual_plus = ComputeResidualVector(state_plus, last_pose, last_velocity, pose, velocity);
		
		// std::cout << " residual_plus : \n" << residual_plus << std::endl;

		Eigen::Matrix<double, 10, 1> residual_minus = ComputeResidualVector(state_minus, last_pose, last_velocity, pose, velocity);
		Eigen::Matrix<double, 10, 1> delta_residual = residual_plus - residual_minus;
		
		// std::cout << " delta_residual : \n" << delta_residual << std::endl;

		Eigen::Matrix<double, 10, 1> scalar_residual = scalar * delta_residual;

		jacobian.block<10, 1>(0, i) = scalar_residual;
	}

	// std::cout << " jacobian : \n" << jacobian << std::endl;

	jacobianXi = jacobian.block<5,8>(0,0);
    jacobianXij = jacobian.block<5,5>(0,8);
    jacobianXji = jacobian.block<5,8>(5,0);
    jacobianXj = jacobian.block<5,5>(5,8);
}


void PoseObserve::ComputeResidual(const State& state, const Eigen::Vector3d &pose, const Eigen::Vector2d &velocity, Eigen::Matrix<double, 5, 1> &residual) {
	Eigen::Vector2d res_p = state.p - pose.block<2,1>(0,0);
	double res_theta = math_utils::normalize_angle(state.theta - pose[2]);
	// double res_theta = math_utils::theta_normalize(state.theta - pose[2]);

	Eigen::Vector2d res_v = state.v - velocity;

	residual.block<2,1>(0,0) = res_p;
	residual.block<2,1>(2,0) = res_v;
	residual(4,0) = res_theta;
}
